import base64
import io
import os
import sys

import gradio as gr

from sdk.httpclient import CTClientBuilder, CTClient

base_url = "https://ai-global.ctapi.ctyun.cn"

client: CTClient


def _get_image_base64(img):
    """
    返回图片的base64大小（单位MB）和内容
    """
    byte_arr = io.BytesIO()
    img.save(byte_arr, format='PNG')
    byte_arr = byte_arr.getvalue()
    base64_str = base64.b64encode(byte_arr).decode('utf-8')
    base64_mb_size = sys.getsizeof(base64_str) / (1024 * 1024)
    return base64_mb_size, base64_str


def ocr(img):
    base64_size, base64_str = _get_image_base64(img)
    if base64_size > 2.0:
        return {"errMsg": f"上传的图片进行base64编码后大小为 {base64_size:.2f} MB，超过 2.00 MB"}, None

    url = f"{base_url}/v1/aiop/api/2f3p1pnxpqm8/ocrdetect/ocr/v1/image.json"
    response = client.request("POST", url, body={"data": [base64_str]})
    response.raise_for_status()
    response_json = response.json()
    try:
        result = response_json['returnObj']
        return result
    except Exception as e:
        print(f"error: {e}")
        return None, response_json


def idcard(img):
    base64_size, base64_str = _get_image_base64(img)
    if base64_size > 2.0:
        return {"errMsg": f"上传的图片进行base64编码后大小为 {base64_size:.2f} MB，超过 2.00 MB"}, None

    url = f"{base_url}/v1/aiop/api/2f3os7qq79xc/IdentityCard/ocr/v1/idcard.json"
    response = client.request("POST", url, body={"data": [base64_str]})
    response.raise_for_status()
    response_json = response.json()
    try:
        result = response_json['returnObj']
        return result
    except Exception as e:
        print(f"error: {e}")
        return None, response_json


def plate_license(img):
    base64_size, base64_str = _get_image_base64(img)
    if base64_size > 2.0:
        return {"errMsg": f"上传的图片进行base64编码后大小为 {base64_size:.2f} MB，超过 2.00 MB"}, None

    url = f"{base_url}/v1/aiop/api/2gt54ed8660w/driven-plate-ocr/platelicense.json"
    response = client.request("POST", url, body={"data": [base64_str]})
    response.raise_for_status()
    response_json = response.json()
    try:
        result = response_json['returnObj']
        return result
    except Exception as e:
        print(f"error: {e}")
        return None, response_json


def business_license(img):
    base64_size, base64_str = _get_image_base64(img)
    if base64_size > 2.0:
        return {"errMsg": f"上传的图片进行base64编码后大小为 {base64_size:.2f} MB，超过 2.00 MB"}, None

    url = f"{base_url}/v1/aiop/api/2k8jo3ghhjwg/businesslicense/ocr/v1/BusinessLicense.json"
    response = client.request("POST", url, body={"Action": "BusinessCardOCR", "ImageData": base64_str})
    response.raise_for_status()
    response_json = response.json()
    try:
        result = response_json['returnObj']
        return result
    except Exception as e:
        print(f"error: {e}")
        return None, response_json


product_intro = "进一步了解天翼云印刷文字识别产品：<a href='https://www.ctyun.cn/products/printedwordocr'>https://www.ctyun.cn/products/printedwordocr</a>"
gr_config = {
    "ocr": gr.Interface(fn=ocr,
                                inputs=gr.Image(type="pil", label="上传图像"),
                                outputs=[
                                    gr.Textbox(label="识别结果")
                                ],
                                examples=[],
                                title="印刷文字识别 - 通用型OCR",
                                description="针对图片中的文字，进行OCR检测，返回检测到的文字内容及坐标信息。",
                                article=product_intro),
    "idcard": gr.Interface(fn=idcard,
                                inputs=gr.Image(type="pil", label="上传图像"),
                                outputs=[
                                    gr.Textbox(label="识别结果")
                                ],
                                examples=[],
                                title="印刷文字识别 - 身份证识别",
                                description="针对图片中的身份证，进行OCR检测，返回检测到的姓名、身份证号码等信息。",
                                article=product_intro),
    "plate_license": gr.Interface(fn=plate_license,
                                inputs=gr.Image(type="pil", label="上传图像"),
                                outputs=[
                                    gr.Textbox(label="识别结果")
                                ],
                                examples=[],
                                title="印刷文字识别 - 车牌识别",
                                description="对图片中的车牌（仅限中国大陆境内的蓝牌、黄牌（单层）、新能源绿牌），进行OCR检测，返回检测到的车牌内容及车牌位置坐标。",
                                article=product_intro),
    "business_license": gr.Interface(fn=business_license,
                                  inputs=gr.Image(type="pil", label="上传图像"),
                                  outputs=[
                                      gr.Textbox(label="识别结果")
                                  ],
                                  examples=[],
                                  title="印刷文字识别 - 营业执照识别",
                                  description="针对营业执照图片，进行OCR检测，返回检测到的统一社会信用代码，注册号，名称，类型，住所，法定代表人，注册资金，成立日期，营业期限，实收资本。",
                                  article=product_intro),
}


def get_not_empty_env(key):
    value = os.getenv(key, "").strip()
    if value == "":
        raise Exception(f"env {key} is not set or empty")
    return value


def run():
    ctyun_ak = get_not_empty_env("ext_cf_ctyun_ak")
    ctyun_sk = get_not_empty_env("ext_cf_ctyun_sk")
    ctyun_ai_app_key = get_not_empty_env("ext_cf_ctyun_ai_app_key")
    app_type = get_not_empty_env("ext_cf_app_type")

    global client
    client = CTClientBuilder().with_ak(ctyun_ak).with_sk(ctyun_sk).with_ai_app_key(ctyun_ai_app_key).build()

    gr_config[app_type].launch(server_name="0.0.0.0", server_port=9000)


if __name__ == '__main__':
    run()
